package com.ybchen.thread;

import com.ybchen.content.TraceContextHolder;

import java.util.concurrent.*;

/**
 * @description: traceId线程池
 * @author: alex
 * @create: 2024-10-20 20:52
 */
public class TraceIdThreadPoolExecutor extends ThreadPoolExecutor {

    public TraceIdThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit,
                                     java.util.concurrent.BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    public TraceIdThreadPoolExecutor(int corePoolSize,
                                     int maximumPoolSize,
                                     long keepAliveTime,
                                     TimeUnit unit,
                                     BlockingQueue<Runnable> workQueue,
                                     String prefix,
                                     RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, new NamedThreadFactory(prefix), handler);
    }

    @Override
    protected void beforeExecute(Thread t, Runnable r) {
        super.beforeExecute(t, r);
        TraceContextHolder.setTraceId();
    }

    @Override
    protected void afterExecute(Runnable r, Throwable t) {
        try {
            super.afterExecute(r, t);
        } finally {
            TraceContextHolder.clear();
        }
    }
}
